{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJcYs_ERTnnI"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HMUDt0CiUJk9"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "77z2OchJTk0l"
      },
      "source": [
        "# Migration examples: Multi-worker training with CPU/GPU\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/migrate/multi_worker_cpu_gpu_training\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/multi_worker_cpu_gpu_training.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/multi_worker_cpu_gpu_training.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/multi_worker_cpu_gpu_training.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "meUTrR4I6m1C"
      },
      "source": [
        "Multi-worker distributed training has been traditionally done in TF1 via `tf.estimator.train_and_evaluate` API with a `tf.estimator.Estimator`. In TF2, it is recommended you write models with metrics, losses, and optimizers in Keras, and distribute them across multiple workers with `Strategy`s found in `tf.distribute`. When it comes to multi-worker training with CPU/GPU, `tf.distribute.experimental.ParameterServerStrategy` and `tf.distribute.MultiWorkerMirroredStrategy` are the strategies of choice, and this guide demonstrates how TF1 distributed training code can migrate to TF2.\n",
        "\n",
        "For more information about other `Strategy`s including when to use what, please refer to [Distributed Training with Tensorflow guide](https://www.tensorflow.org/guide/distributed_training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdZSoIXEbhg-"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Start with a couple of necessary install and TensorFlow imports:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iE0vSfMXumKI"
      },
      "outputs": [],
      "source": [
        "# Temporarily install tf-nightly as the notebook depends on symbols in 2.6.\n",
        "!pip uninstall -q -y tensorflow keras\n",
        "!pip install -q tf-nightly\n",
        "# Utility needed for this demonstration\n",
        "!pip install portpicker\n",
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jsm9Rxx7s1OZ"
      },
      "source": [
        "and prepare some simple data for demonstration:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m7rnGxsXtDkV"
      },
      "outputs": [],
      "source": [
        "features = [[1., 1.5], [2., 2.5], [3., 3.5]]\n",
        "labels = [[0.3], [0.5], [0.7]]\n",
        "eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]\n",
        "eval_labels = [[0.8], [0.9], [1.]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T2uaw9QaDM_X"
      },
      "source": [
        "The environment variable `TF_CONFIG` is required in TF1 to specify the cluster, and the tasks' addresses. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4OUzwoQgXgkG"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "\n",
        "tf_config = {\n",
        "    'cluster': {\n",
        "        'chief': ['localhost:11111'],\n",
        "        'worker': ['localhost:12345', 'localhost:23456', 'localhost:21212'],\n",
        "        'ps': ['localhost:12121', 'localhost:13131'],\n",
        "    },\n",
        "    'task': {'type': 'chief', 'index': 0}\n",
        "}\n",
        "\n",
        "os.environ['TF_CONFIG'] = json.dumps(tf_config)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PbeoSbbmDdc0"
      },
      "source": [
        "Unfortunately, since TF1 multi-worker training with `Estimator` requires multiple  clients, which would be especially tricky to be done here in a notebook, you will make the notebook be runnable without a `TF_CONFIG` so it falls back to local training. Note that regular TF1 multi-worker training does use the `TF_CONFIG` without the following `del` this notebook performs here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AHuynAR5D8sU"
      },
      "outputs": [],
      "source": [
        "del os.environ['TF_CONFIG']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4uXff1BEssdE"
      },
      "source": [
        "### TF1: Using `tf.estimator.train_and_evaluate`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MpyINdiLEN3c"
      },
      "source": [
        "This code snippet demonstrates the canonical workflow of multi-worker training in TF1: you will use a `tf.estimator.Estimator`, a `tf.estimator.TrainSpec`, a `tf.estimator.EvalSpec`, and the `tf.estimator.train_and_evaluate` API to distribute the training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lqe9obf7suIj"
      },
      "outputs": [],
      "source": [
        "def _input_fn():\n",
        "  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)\n",
        "\n",
        "def _eval_input_fn():\n",
        "  return tf1.data.Dataset.from_tensor_slices(\n",
        "      (eval_features, eval_labels)).batch(1)\n",
        "\n",
        "def _model_fn(features, labels, mode):\n",
        "  logits = tf1.layers.Dense(1)(features)\n",
        "  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)\n",
        "  optimizer = tf1.train.AdagradOptimizer(0.05)\n",
        "  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())\n",
        "  return tf1.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)\n",
        "\n",
        "estimator = tf1.estimator.Estimator(model_fn=_model_fn)\n",
        "train_spec = tf1.estimator.TrainSpec(input_fn=_input_fn)\n",
        "eval_spec = tf1.estimator.EvalSpec(input_fn=_eval_input_fn)\n",
        "tf1.estimator.train_and_evaluate(estimator, train_spec, eval_spec)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEmzBjfnsxwT"
      },
      "source": [
        "### TF2: Keras training API with `tf.distribute.Strategy`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Syb66qsbEp1x"
      },
      "source": [
        "In TF2, computation is distributed via `tf.distribute.Strategy`s. In this example, you will see two `Strategy`s: `tf.distribute.experimental.ParameterServerStrategy` and `tf.distribute.MultiWorkerMirroredStrategy`, as these two strategies have been designed for CPU/GPU training with multiple workers.\n",
        "\n",
        "`ParameterServerStrategy` employs a coordinator which make it more friendly with the environment this notebook is in, and you will be using some utilities here to set up the supporting elements essential for a runnable experience here: you will create a so-called \"in-process cluster\", where threads are used to simulate the parameter servers (PS) and workers. \n",
        "\n",
        "For more information about parameter server training, please refer to [Parameter Server Training guide](https://www.tensorflow.org/tutorials/distribute/parameter_server_training) \n",
        "\n",
        "Here, you will first define the `TF_CONFIG` environment variable with a `tf.distribute.cluster_resolver.TFConfigClusterResolver` to provide the cluster information.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rp-gFY0H5rF-"
      },
      "outputs": [],
      "source": [
        "# Find ports that are available for chief (the coordinator), workers, and PS.\n",
        "import portpicker\n",
        "chief_port = portpicker.pick_unused_port()\n",
        "worker_ports = [portpicker.pick_unused_port() for _ in range(3)]\n",
        "ps_ports = [portpicker.pick_unused_port() for _ in range(2)]\n",
        "\n",
        "# Dump the cluster information to `TF_CONFIG`.\n",
        "tf_config = {\n",
        "    'cluster': {\n",
        "        'chief': [\"localhost:%s\" % chief_port],\n",
        "        'worker': [\"localhost:%s\" % port for port in worker_ports],\n",
        "        'ps':  [\"localhost:%s\" % port for port in ps_ports],\n",
        "    },\n",
        "    'task': {'type': 'chief', 'index': 0}\n",
        "}\n",
        "os.environ['TF_CONFIG'] = json.dumps(tf_config)\n",
        "\n",
        "# Use a cluster resolver to bridge the info to the strategy created below.\n",
        "cluster_resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o_8uVvJb6dqq"
      },
      "source": [
        "Then, you will create those servers for the workers and PS one-by-one, and end up creating a `tf.distribute.experimental.ParameterServerStrategy`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZJopinmG6b2z"
      },
      "outputs": [],
      "source": [
        "# Workers need some inter_ops threads to work properly.\n",
        "worker_config = tf.compat.v1.ConfigProto()\n",
        "worker_config.inter_op_parallelism_threads = 4\n",
        "\n",
        "for i in range(3):\n",
        "  tf.distribute.Server(\n",
        "      cluster_resolver.cluster_spec(),\n",
        "      job_name=\"worker\",\n",
        "      task_index=i,\n",
        "      config=worker_config)\n",
        "\n",
        "for i in range(2):\n",
        "  tf.distribute.Server(\n",
        "      cluster_resolver.cluster_spec(),\n",
        "      job_name=\"ps\",\n",
        "      task_index=i)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IpfCcF0g6Ao8"
      },
      "source": [
        "Note that in real distributed training, instead of starting all the `tf.distribute.Server`s on the coordinator, you will use multiple machines, and the ones that are designated as worker and PS will each run a `Server`. \n",
        "\n",
        "With everything ready, you will now create the `ParameterServerStrategy` object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t45iQeBT7Us_"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "diNsps1MGRS6"
      },
      "source": [
        "Once you define a strategy object, you will create the model, optimizer, and other variables, as well as calling `Model.compile` within the `strategy.scope` API in order to distribute the training. Please refer to [`Strategy.scope` guide](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#scope) for more information about this API. If you prefer to define the details including the actual forward/backward passes, please see [parameter server training guide](https://www.tensorflow.org/tutorials/distribute/parameter_server_training) for more information.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "atVciNgPs0fw"
      },
      "outputs": [],
      "source": [
        "dataset = tf.data.Dataset.from_tensor_slices(\n",
        "      (features, labels)).shuffle(10).repeat().batch(64)\n",
        "\n",
        "eval_dataset = tf.data.Dataset.from_tensor_slices(\n",
        "      (eval_features, eval_labels)).repeat().batch(1)\n",
        "\n",
        "with strategy.scope():\n",
        "  model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])\n",
        "  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)\n",
        "  model.compile(optimizer, \"mse\")\n",
        "\n",
        "model.fit(dataset, epochs=5, steps_per_epoch=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "akZ0aaaS1vA9"
      },
      "outputs": [],
      "source": [
        "model.evaluate(eval_dataset, steps=10, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pXbS71XmMSoO"
      },
      "source": [
        "> **Partitioners**\n",
        "\n",
        "> ParameterServerStrategy supports variable partitioning and offers same partitioners as TF1, with less confusing names:\n",
        "- variable_axis_size_partitioner -> MaxSizePartitioner (Partitioner that keeps shards under a maximum size)\n",
        "- min_max_variable_partitioner -> MinSizePartitioner (Partitioner that allocates a minimum size per shard)\n",
        "- fixed_size_partitioner -> FixedShardsPartitioner (Partitioner that allocates a fixed number of shards)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ig0-uCUbGprd"
      },
      "source": [
        "Alternatively, a `MultiWorkerMirroredStrategy` object can be used:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xHXP8bOBGtXL"
      },
      "outputs": [],
      "source": [
        "del os.environ['TF_CONFIG']  # To clean up the `TF_CONFIG` used for PS strategy.\n",
        "strategy = tf.distribute.MultiWorkerMirroredStrategy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tOsmqefTGwUf"
      },
      "source": [
        "And you can replace the strategy used above with this `MultiWorkerMirroredStrategy` object to perform a training with this strategy. \n",
        "\n",
        "Unfortunately, as with estimator, since `MultiWorkerMirroredStrategy` is a multi-client strategy, there is not an easy way to run distributed training in this notebook, so replacing code above with this strategy ends up running things locally. For an experience running `MultiWorkerMirroredStrategy` distributedly in a colab, please refer to [Multi-worker training with Keras](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras) guide."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "multi_worker_cpu_gpu_training.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
